#pragma once

#include "../core/i_struct.h"
#include "../core/i_rand.h"
#include "../core/i_blas.h"

#include <algorithm>

namespace idl
{
  /*Compute the number of trials for Ransac. Number of trials is chosen sufficiently high to 
    ensure with a probability "confidence" that at least one of the random samples of data is
    free from outliers. "inlierprob" is the probability that any selected data point is an inlier*/
  inline int i_ransac_trials(int sample_size, double confidence, double inlierprob)
  {
    return inlierprob > 0.0 ? i_round(i_div(i_log(1.0 - confidence), i_log(1.0 - i_pow(inlierprob, sample_size)))) : Constant<int>::MAX_VAL();
  }

  /*Using Ransac to fit a model to data set which contains outliers
    The function needs 2n entries of scratch space in inliers*/
  template <typename T, 
            int l,/*data element size*/ 
            int k,/*# of parameters of the model*/ 
            int s,/*min # of sample size to fit a model*/
            void(*hypogenFunc)(const T* samples, T* model),/*the minimal case solver (hypothesis generator) to fit a model from s samples*/
            void(*costFunc)(const T* model, const T* data, int n, int &nr_liner, int *inliers, T &cost, T error_tol),/*the cost function to validate a model*/
            void(*refitFunc)(T* data, int *inliers, T* model, int n, int nr_liner)>
  bool robustFitRansac(T* data, /*data is not const since it maybe changed by the refitFunc (e.g. line fitting functions) and random shuffle function. If turn off "refitting" and "random shuffle" options the data array won't be destroyed*/
                       int n, /*total # of samples*/
                       T model[k], /*the model, which contains k parameters*/
                       int &consensus_size, /*return the number of inliers in data*/
                       int *inliers, /*scratch space of length 2*n to store the indices of inliers*/
                       T error_tol,
                       bool re_est_model_w_inliers = false, /*set this value to true to use all inliers to re-estimate a model*/
                       bool adaptive_trial_count = false, /*since "inlierprob" is usually unknown, set this value to true will use an adaptive algorithm for determining the number of Ransac trials*/
                       double confidence = 0.99,
                       double inlierprob = 0.5,/*initial guess for the inliers probability (usually worst case)*/
                       int min_nr_inliers = s,
                       bool random_shuffle_inputs = false)   
  {
    int indices[s];
    T samples[l*s];
    T tmp_model[k];
    T cost = Constant<T>::MAX_VAL();
    T best_cost = Constant<T>::MAX_VAL();
    
    if (n < min_nr_inliers)
    {
      return false;
    }

    double actual_inlierprob = 0.0, tmp_inlierprob;
    int nr_trials = i_ransac_trials(s, confidence, inlierprob);
    int nr_inliers = 0;
    int rseed = I_DEFAULT_SEED;
    int sample_count = 0;
    consensus_size = 0; /*initialize the size of the consensus set to zero*/
    i_zero(model, k); /*initialize the model with zeros*/

    if (random_shuffle_inputs)
    {
      i_randomized_shuffle(data, n, l, rseed);
    }

    while (nr_trials > sample_count)
    {
      /*generate random indices*/
      i_random_sample(indices, s, n, rseed);
      /*prepare data for model fitting*/
      for (int i = 0; i < s; ++i)
      {
        i_copy(data + indices[i]*l, samples+i*l, l);
      }
      /*estimate model*/
      hypogenFunc(samples, tmp_model);
      /*validate model*/
      costFunc(tmp_model, data, n, nr_inliers, inliers+n/*use the second half of scratch space*/, cost, error_tol);
      if ((nr_inliers > consensus_size) || (nr_inliers == consensus_size && cost < best_cost)) /*found a better model*/
      {
        consensus_size = nr_inliers;
        best_cost = cost;
        i_copy(tmp_model, model, k);
        i_copy(inliers + n, inliers, consensus_size);/*record inlier indices*/
        if (adaptive_trial_count)
        {
          tmp_inlierprob = i_div((double)consensus_size, n);
          if (tmp_inlierprob > actual_inlierprob)
          {
            actual_inlierprob = tmp_inlierprob;
            nr_trials = i_ransac_trials(s, confidence, actual_inlierprob);
          }
        }
      }
      sample_count++;
    }
    bool succeeded = consensus_size >= min_nr_inliers; 
    if (succeeded && re_est_model_w_inliers && refitFunc != NULL)
    {
      refitFunc(data, inliers, model, n, consensus_size);
    }
    return succeeded;
  }

  /*Using Ransac to fit a model to data set which contains outliers
  The function needs 2n entries of scratch space in inliers*/
  template <typename T,
            int l,/*data element size of x*/
            int lp,/*data element size of xp*/
            int k,/*# of parameters of the model*/
            int s,/*min # of sample size to fit a model*/
            void(*hypogenFunc)(const T* x, const T* xp, T* model),/*the minimal case solver (hypothesis generator) to fit a model from s samples*/
            void(*costFunc)(const T* model, const T* x, const T* xp, int n, int &nr_liner, int *inliers, T &cost, T error_tol),/*the cost function to validate a model*/
            void(*refitFunc)(T* x, T* xp, int *inliers, T* model, int n, int nr_liner)>
            bool robustBinaryFitRansac(T* x, /*data is not const since it maybe changed by the refitFunc (e.g. line fitting functions) and random shuffle function. If turn off "refitting" and "random shuffle" options the data array won't be destroyed*/
                             T* xp,
                             int n, /*total # of samples*/
                             T model[k], /*the model, which contains k parameters*/
                             int &consensus_size, /*return the number of inliers in data*/
                             int *inliers, /*scratch space of length 2*n to store the indices of inliers*/
                             T error_tol,
                             bool re_est_model_w_inliers = false, /*set this value to true to use all inliers to re-estimate a model*/
                             bool adaptive_trial_count = false, /*since "inlierprob" is usually unknown, set this value to true will use an adaptive algorithm for determining the number of Ransac trials*/
                             double confidence = 0.99,
                             double inlierprob = 0.5,/*initial guess for the inliers probability (usually worst case)*/
                             int min_nr_inliers = s,
                             bool random_shuffle_inputs = false)
  {
    int indices[s];
    T samples_x[l*s];
    T samples_xp[lp*s];
    T tmp_model[k];
    T cost = Constant<T>::MAX_VAL();
    T best_cost = Constant<T>::MAX_VAL();

    if (n < min_nr_inliers)
    {
      return false;
    }

    double actual_inlierprob = 0.0, tmp_inlierprob;
    int nr_trials = i_ransac_trials(s, confidence, inlierprob);
   
    int nr_inliers = 0;
    int rseed = I_DEFAULT_SEED;
    int sample_count = 0;
    int i, idxl, idxlp, il, ilp;
    consensus_size = 0; /*initialize the size of the consensus set to zero*/
    i_zero(model, k); /*initialize the model with zeros*/

    if (random_shuffle_inputs)
    {
      i_randomized_shuffle(x, xp, n, l, lp, rseed);
    }

    while (nr_trials > sample_count)
    {
      /*generate random indices*/
      i_random_sample(indices, s, n, rseed);
      /*prepare data for model fitting*/
      for (i = 0; i < s; ++i)
      {
        idxl  = indices[i] * l;
        idxlp = indices[i] * lp;
        il = i*l;
        ilp = i*lp;
        i_copy(x  + idxl,  samples_x  + il,  l);
        i_copy(xp + idxlp, samples_xp + ilp, lp);
      }
      
      /*estimate model*/
      hypogenFunc(samples_x, samples_xp, tmp_model);
      
      /*validate model*/
      costFunc(tmp_model, x, xp, n, nr_inliers, inliers + n/*use the second half of scratch space - that's why the size of "inliers" is 2*n*/, cost, error_tol);
      if ((nr_inliers > consensus_size) || (nr_inliers == consensus_size && cost < best_cost)) /*found a better model*/
      {
        consensus_size = nr_inliers;
        best_cost = cost;
        i_copy(tmp_model, model, k);
        i_copy(inliers + n, inliers, consensus_size);/*record inlier indices*/
        if (adaptive_trial_count)
        {
          tmp_inlierprob = i_div((double)consensus_size, n);
          if (tmp_inlierprob > actual_inlierprob)
          {
            actual_inlierprob = tmp_inlierprob;
            nr_trials = i_ransac_trials(s, confidence, actual_inlierprob);
          }
        }
      }
      sample_count++; 
    }
    bool succeeded = consensus_size >= min_nr_inliers;/*(i_div((double)consensus_size, n) >= inlierprob) -> too restricted*/;
    if (succeeded && re_est_model_w_inliers && refitFunc != NULL)
    {
      refitFunc(x, xp, inliers, model, n, consensus_size);
    }
    return succeeded;
  }

  /*Using David Nister's Preemptive Ransac - "Preemptive RANSAC for Live Structure and Motion Estimation, ICCV 2003"
  to fit a model to data set which contains outliers
  The function needs n entries of scratch space in inliers - not 2n*/
  template <typename T,
    int l,/*data element size of x*/
    int lp,/*data element size of xp*/
    int k,/*# of parameters of the model*/
    int s,/*min # of sample size to fit a model*/
    int m, /*# of total hypothesis*/
    int b, /*block size - suggested value 0.2*m, b must be less or equal to m*/
    void(*hypogenFunc)(const T* x, const T* xp, T* model),/*the minimal case solver (hypothesis generator) to fit a model from s samples*/
    void(*costFunc)(const T* model, const T* x, const T* xp, T &cost)/*the cost function to validate a model, only for one data point*/>
    bool robustBinaryFitPreemptiveRansac(T* x, /*data is not const since it will be changed by the random shuffle*/
    T* xp,
    int n, /*total # of samples*/
    T model[k], /*the model, which contains k parameters*/
    int &consensus_size, /*return the number of inliers in data*/
    int *inliers, /*scratch space of length n to store the indices of inliers*/
    T error_tol,
    int min_nr_inliers = s)
  {
      int i, j, fi, fi_half, idxl, idxlp, ioverb, indices[s], best_index, best_num;
      int rseed = I_DEFAULT_SEED;
      T samples_x[l*s];
      T samples_xp[lp*s];
      T models[k * m];
      T *px, *pxp, cost;
      Pair<int, int> counts[m];
      consensus_size = 0; /*initialize the size of the consensus set to zero*/
      i_zero(model, k); /*initialize the model with zeros*/

      if (b > m || n < min_nr_inliers)
      {
        return false;
      }

      /*random shuffle the inputs*/
      i_randomized_shuffle(x, xp, n, l, lp, rseed);
      
      /*hypothesis generation*/
      for (i = 0; i < m; ++i)
      {
        i_random_sample(indices, s, n, rseed);
        
        /*prepare data for model fitting*/
        for (j = 0; j < s; ++j)
        {
          idxl  = indices[j] * l;
          idxlp = indices[j] * lp;
          i_copy(x + idxl,   samples_x  + j*l,  l);
          i_copy(xp + idxlp, samples_xp + j*lp, lp);
        }

        /*generate the ith model hypothesis*/
        hypogenFunc(samples_x, samples_xp, models + k*i);

        /*initialize the count buffer*/
        counts[i].first  = i; /*initialize to index id*/
        counts[i].second = 0; /*initialize to count 0*/
      }

      /*test hypothesis*/
      i = ioverb = 0;
      fi = m;

      while (i < n && fi > 1)
      {
        px  = x  + i*l;
        pxp = xp + i*lp;

        for (j = 0; j < fi; ++j)
        {
          costFunc(models + (counts[j].first * k), px, pxp, cost);
          if (cost < error_tol)
          {
            counts[j].second++;
          }
        }
        
        ++i;

        if (ioverb < i / b)
        {
          ioverb++;
          fi_half = (fi >> 1); /*shift right -> divided by 2*/
          std::partial_sort(counts, counts + fi_half, counts + fi, i_larger_pair_second_element<int, int>);
          fi = fi_half;
        }
      }

      best_index = counts[0].first;
      best_num   = counts[0].second;

      if (n == i && fi > 1)
      {
        for (j = 1; j < fi; ++j)
        {
          if (counts[j].second > best_num)
          {
            best_index = counts[j].first;
            best_num = counts[j].second;
          }
        }
      }

      i_copy(models + (best_index*k), model, k); /*after the previous iteration, the first model is the best solution with max # of supports - partial_sort guarantees sorted first-half*/

      /*reassemble inliers*/
      px  = x;
      pxp = xp;
      for (j = 0; j < n; ++j)
      {
        costFunc(model, px, pxp, cost);
        if (cost < error_tol)
        {
          inliers[consensus_size++] = j;
        }
        px += l;
        pxp += lp;
      }

      bool succeeded = consensus_size >= min_nr_inliers;
      return succeeded;
  }

}/* namespace idl */